//! Tracker for local variables.

use std::collections::HashSet;

use thiserror::Error;

use stencil_tysys::{EnvExpr, LocalRef, ValueType};

/// Errors involved in tracking variables.
#[derive(Debug, Error)]
pub enum VarError {
    #[error("local {0:?} has not been assigned a value")]
    LocalUnassigned(LocalRef),

    #[error("tried to reassign to {0:?} with wrong type (is {1:?}, tried {2:?}")]
    TypeConflict(EnvExpr, ValueType, ValueType),
}

/// Tracks locals that are in use and not.
pub struct LocalVarTracker {
    /// Locals' types.
    locals: Vec<ValueType>,

    /// Locals that have values assigned.
    locals_assigned: HashSet<LocalRef>,

    /// Locals that have been *reassigned*.
    locals_mutated: HashSet<LocalRef>,
}

impl LocalVarTracker {
    pub fn new() -> Self {
        Self {
            locals: Vec::new(),
            locals_assigned: HashSet::new(),
            locals_mutated: HashSet::new(),
        }
    }

    /// Returns a (currently-invalid) reference to the "next" local to be produced.
    fn next_local(&self) -> LocalRef {
        let i = self.locals.len();
        LocalRef::from(i as u16)
    }

    /// Declare's a new local, without assigning it a value.
    pub fn declare_local(&mut self, ty: ValueType) -> LocalRef {
        let lr = self.next_local();
        self.locals.push(ty);
        lr
    }

    /// Assigns to a new local, returning a ref.
    pub fn assign_new_local(&mut self, ty: ValueType) -> LocalRef {
        let lr = self.next_local();
        self.locals.push(ty);
        self.locals_assigned.insert(lr);
        lr
    }

    /// Assign a local, returning if it was freshly or dirtyly.
    pub fn reassign_local(&mut self, lr: LocalRef, ty: ValueType) -> Result<bool, VarError> {
        if let Some(cur_ty) = self.locals.get(lr.inner() as usize) {
            if cur_ty != &ty {
                return Err(VarError::TypeConflict(
                    EnvExpr::Local(lr),
                    cur_ty.clone(),
                    ty,
                ));
            }
        }

        // Update the fields.
        // TODO verify if I have this the right way around
        let is_already_assigned = self.locals_assigned.insert(lr);
        if !is_already_assigned {
            self.locals_mutated.insert(lr);
        }

        Ok(is_already_assigned)
    }

    /// Gets the type of a local variable.
    pub fn get_local_ty(&self, lr: LocalRef) -> Option<&ValueType> {
        self.locals.get(lr.inner() as usize)
    }
}
